package websocket

import (
	"context"
	"github.com/gogf/gf/v2/net/ghttp"
	pkgErr "github.com/pkg/errors"
	"sync"
	"time"
	"websocket/internal/consts"
)

type wAssemble struct {
	client       map[*ghttp.WebSocket]*client                    //非并发安全 操作时需要手动上锁
	uniqueClient map[consts.WebsocketUniqueType]*ghttp.WebSocket //非并发安全 操作时需要手动上锁
	mux          *sync.RWMutex                                   //读写锁 - 可改写成对unique的内存互斥锁
}

var (
	insAssemble = wAssemble{
		client:       make(map[*ghttp.WebSocket]*client),
		uniqueClient: make(map[consts.WebsocketUniqueType]*ghttp.WebSocket),
		mux:          new(sync.RWMutex),
	}
)

func Assemble() *wAssemble {
	return &insAssemble
}

func (w *wAssemble) SetClient(ctx context.Context, unique consts.WebsocketUniqueType, conn *ghttp.WebSocket, c *client) error {
	w.mux.Lock()
	defer w.mux.Unlock()
	if _, ok := w.client[conn]; ok {
		return pkgErr.New("client repeat setting")
	}
	w.client[conn] = c
	w.uniqueClient[unique] = conn
	return nil
}

func (w *wAssemble) GetClient(ctx context.Context, conn *ghttp.WebSocket) (*client, error) {
	w.mux.Lock()
	defer w.mux.Unlock()
	c, ok := w.client[conn]
	if !ok {
		return nil, pkgErr.New("client is not set")
	}
	return c, nil
}

func (w *wAssemble) GetClientByUnique(ctx context.Context, unique consts.WebsocketUniqueType) (*client, error) {
	w.mux.Lock()
	defer w.mux.Unlock()
	conn, ok := w.uniqueClient[unique]
	if !ok {
		return nil, pkgErr.New("unique is not set")
	}
	c, ok := w.client[conn]
	if !ok {
		return nil, pkgErr.New("client is not set")
	}
	return c, nil
}

func (w *wAssemble) RemoveClient(ctx context.Context, unique consts.WebsocketUniqueType, force bool) {
	w.mux.Lock()
	defer w.mux.Unlock()

	conn, ok := w.uniqueClient[unique]
	if !ok {
		return
	}
	delete(w.uniqueClient, unique)
	c, ok := w.client[conn]
	if !ok {
		return
	}
	c.Closed = true
	c.Heartbeat().Close()
	if !force {
		//重试次数
		retry := 0
		for {
			if c.WriteChNum == 0 {
				break
			}
			retry++
			time.Sleep(500 * time.Millisecond)
			if retry > 20 {
				//达到最大重试次数直接退出 防止死循环
				break
			}
		}
	}
	c.ctxClose()
	close(c.writeCh)
	delete(w.client, conn)
	_ = conn.Close()
	return
}

func (w *wAssemble) GetClients(ctx context.Context) map[consts.WebsocketUniqueType]*client {
	clients := make(map[consts.WebsocketUniqueType]*client)
	for unique, conn := range w.uniqueClient {
		clients[unique] = w.client[conn]
	}
	return clients
}
